Description

HuggingFaceGenerationAdapter.prepare_inputs_for_generation builds model_inputs with

"tensor_capture_hook": tensor_capture_hook,

but never extracts tensor_capture_hook from **kwargs, while input_capture_hook is extracted on the line immediately above it. Any call path that reaches this method — including adapter.generate() on a model whose forward() accepts a tensor_capture_hook kwarg (e.g. NeuronBaseForImageToText subclasses) — raises NameError: name 'tensor_capture_hook' is not defined.

This PR adds the missing one-line extraction, mirroring how input_capture_hook is handled.

Reproduction

Any call to adapter.generate(...) where the underlying model declares tensor_capture_hook in its forward signature triggers the bug. Minimal example:

from neuronx_distributed_inference.utils.hf_adapter import HuggingFaceGenerationAdapter
# adapter wraps any NeuronBaseForImageToText-based model (e.g. Qwen2-VL, Qwen3-VL)
adapter.generate(input_ids=..., attention_mask=..., max_new_tokens=8)
# -> NameError: name 'tensor_capture_hook' is not defined

Fix

src/neuronx_distributed_inference/utils/hf_adapter.py:

         input_capture_hook = kwargs.get("input_capture_hook", None)
+        tensor_capture_hook = kwargs.get("tensor_capture_hook", None)

Behavior for callers that do not pass tensor_capture_hook is unchanged (it stays None, which is what model_inputs["tensor_capture_hook"] = None already means in the existing code).

Testing

• Manually verified that adapter.generate() on an image-to-text model no longer raises NameError.
• No change to the public API surface; the only change is that a previously undefined local is now extracted from **kwargs with a None default, matching the pattern already used for input_capture_hook and accepted_indices.

Compatibility

• Neuron SDK Version(s): 2.27.x (where the bug was introduced) and newer
• Instance Type(s): Any
• PyTorch Version: Any supported
• Python Version: Any supported

Related

Surfaced while refactoring #135 (Qwen2.5-Omni-7B contrib), whose Talker drives generation via adapter.generate() and therefore hits this path. Once this fix lands, the workaround shim in contrib/models/Qwen2.5-Omni-7B/src/_upstream_compat.py can be removed.

By submitting this PR, I confirm that:

• The code follows best practices and is well-documented